package com.zrc.japktool.apk.merge.dex;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import org.jf.dexlib2.Opcode;
import org.jf.dexlib2.iface.*;
import org.jf.dexlib2.iface.debug.DebugItem;
import org.jf.dexlib2.iface.instruction.Instruction;
import org.jf.dexlib2.iface.instruction.formats.Instruction31i;
import org.jf.dexlib2.iface.reference.FieldReference;
import org.jf.dexlib2.iface.reference.MethodReference;
import org.jf.dexlib2.iface.value.EncodedValue;
import org.jf.dexlib2.iface.value.IntEncodedValue;
import org.jf.dexlib2.immutable.*;
import org.jf.dexlib2.immutable.instruction.ImmutableInstruction;
import org.jf.dexlib2.immutable.instruction.ImmutableInstruction31i;
import org.jf.dexlib2.immutable.reference.ImmutableFieldReference;
import org.jf.dexlib2.immutable.reference.ImmutableMethodReference;
import org.jf.dexlib2.immutable.value.ImmutableIntEncodedValue;
import org.jf.dexlib2.rewriter.Rewriter;
import org.jf.dexlib2.rewriter.Rewriters;
import org.jf.dexlib2.rewriter.RewriterModule;
import org.jf.dexlib2.rewriter.DexRewriter;
import org.jf.dexlib2.rewriter.RewriterUtils;


import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

public class ModifiyDexClassID1 {

    public static void modifiyClassIds(List<ClassDef> targetList, Map<String, String> ids) {
        RewriterModule module = new RewriterModule() {
            // 重写类型引用（类名、父类、接口等）
            @Override
            public Rewriter<String> getTypeRewriter(Rewriters rewriters) {
                return type -> ids.getOrDefault(type, type);
            }

            // 重写字段引用
            @Override
            public Rewriter<FieldReference> getFieldReferenceRewriter(Rewriters rewriters) {
                return fieldRef -> {
                    String newDefiningClass = rewriters.getTypeRewriter().rewrite(fieldRef.getDefiningClass());
                    String newType = rewriters.getTypeRewriter().rewrite(fieldRef.getType());
                    return new ImmutableFieldReference(newDefiningClass, fieldRef.getName(), newType);
                };
            }

            // 重写方法引用
            @Override
            public Rewriter<MethodReference> getMethodReferenceRewriter(Rewriters rewriters) {
                return methodRef -> {
                    String newDefiningClass = rewriters.getTypeRewriter().rewrite(methodRef.getDefiningClass());
                    List<String> newParameters = new ArrayList<>();
                    for (CharSequence param : methodRef.getParameterTypes()) {
                        newParameters.add(rewriters.getTypeRewriter().rewrite(param.toString()));
                    }
                    String newReturnType = rewriters.getTypeRewriter().rewrite(methodRef.getReturnType());
                    return new ImmutableMethodReference(newDefiningClass, methodRef.getName(), newParameters, newReturnType);
                };
            }

            // 重写字段，处理字段初始值
            @Override
            public Rewriter<Field> getFieldRewriter(Rewriters rewriters) {
                return field -> {
                    FieldReference rewrittenFieldRef = rewriters.getFieldReferenceRewriter().rewrite(field);
                    EncodedValue originalValue = field.getInitialValue();// 处理初始值（仅修改需要替换的int类型值，其他类型保持原样）
                    EncodedValue rewrittenValue = originalValue;// 默认保持原始值

                    if (originalValue instanceof IntEncodedValue) {
                        int originalInt = ((IntEncodedValue) originalValue).getValue();
                        String hexKey = String.format("0x%08x", originalInt);
                        if (ids.containsKey(hexKey)) {
                            try {
                                int newValue = Integer.decode(ids.get(hexKey));
                                rewrittenValue = new ImmutableIntEncodedValue(newValue);
                            } catch (NumberFormatException e) {
                                rewrittenValue = originalValue;// 解析失败时保留原始值
                            }
                        }
                    }

                    return new ImmutableField(
                            rewrittenFieldRef.getDefiningClass(),// 重写后的类名（遵循原有重写逻辑）
                            rewrittenFieldRef.getName(),// 重写后的字段名（遵循原有重写逻辑）
                            rewrittenFieldRef.getType(),// 重写后的类型（遵循原有重写逻辑）
                            field.getAccessFlags(),// 保留原始访问标志
                            rewrittenValue,// 仅修改初始值（需要时）
                            field.getAnnotations(),// 保留原始注解（如需重写注解可使用rewriters.getAnnotationRewriter()）
                            field.getHiddenApiRestrictions() // 保留原始隐藏API限制
                    );
                };
            }

            // 重写方法实现，处理指令中的常量值（如const指令）
            @Override
            public Rewriter<MethodImplementation> getMethodImplementationRewriter(Rewriters rewriters) {
                return implementation -> {
                    // 处理 implementation 为 null 的情况（如抽象方法、native方法）
                    if (implementation == null) {
                        return null;
                    }

                    List<ImmutableInstruction> rewrittenInstructions = new ArrayList<>();
                    for (Instruction instruction : implementation.getInstructions()) {
                        // 处理const指令（31i格式，如const v0, 0x7f0b001c）
                        if (instruction.getOpcode() == Opcode.CONST) {
                            Instruction31i constInstruction = (Instruction31i) instruction;
                            int originalValue = constInstruction.getNarrowLiteral();
                            String hexKey = String.format("0x%08x", originalValue);

                            if (ids.containsKey(hexKey)) {
                                try {
                                    int newValue = Integer.decode(ids.get(hexKey));
                                    // 创建新的常量指令
                                    rewrittenInstructions.add(new ImmutableInstruction31i(
                                            Opcode.CONST,
                                            constInstruction.getRegisterA(),
                                            newValue  // 确保newValue是32位有符号整数
                                    ));
                                    continue;
                                } catch (NumberFormatException e) {
                                    // 格式错误则保留原指令
                                }
                            }
                        }
                        // 其他指令直接保留（或转换为不可变对象）
                        rewrittenInstructions.add(ImmutableInstruction.of(instruction));
                    }

                    // 重写try块和调试信息
                    List<? extends TryBlock<? extends ExceptionHandler>> rewrittenTryBlocks = RewriterUtils.rewriteList(rewriters.getTryBlockRewriter(), implementation.getTryBlocks());
                    Iterable<? extends DebugItem> rewrittenDebugItems = RewriterUtils.rewriteIterable(rewriters.getDebugItemRewriter(), implementation.getDebugItems());
                    return new ImmutableMethodImplementation(implementation.getRegisterCount(), rewrittenInstructions, rewrittenTryBlocks, rewrittenDebugItems);
                };
            }

            @Override
            public Rewriter<Method> getMethodRewriter(Rewriters rewriters) {
                return method -> {
                    MethodReference rewrittenMethodRef = rewriters.getMethodReferenceRewriter().rewrite(method);
                    MethodImplementation rewrittenImplementation = rewriters.getMethodImplementationRewriter()
                            .rewrite(method.getImplementation());

                    // 修复1：使用通配符接收参数列表
                    List<? extends MethodParameter> originalParameters = method.getParameters();
                    List<ImmutableMethodParameter> parameters = originalParameters.stream()
                            .map(originalParam -> {
                                String rewrittenType = rewriters.getTypeRewriter().rewrite(originalParam.getType());
                                Set<? extends Annotation> originalAnnotations = originalParam.getAnnotations();

                                // 修复2：使用rewriteSet处理注解集合
                                Set<? extends Annotation> rewrittenAnnotations = RewriterUtils.rewriteSet(rewriters.getAnnotationRewriter(), originalAnnotations);

                                // 转换为ImmutableAnnotation集合
                                ImmutableSet<ImmutableAnnotation> immutableAnnotations = ImmutableAnnotation.immutableSetOf(rewrittenAnnotations);

                                return new ImmutableMethodParameter(rewrittenType, immutableAnnotations, originalParam.getName());
                            })
                            .collect(Collectors.collectingAndThen(
                                    Collectors.toList(),
                                    ImmutableList::copyOf
                            ));

                    return new ImmutableMethod(
                            rewrittenMethodRef.getDefiningClass(),
                            rewrittenMethodRef.getName(),
                            parameters,
                            rewrittenMethodRef.getReturnType(),
                            method.getAccessFlags(),
                            RewriterUtils.rewriteSet(rewriters.getAnnotationRewriter(), method.getAnnotations()),
                            method.getHiddenApiRestrictions(),
                            rewrittenImplementation
                    );
                };
            }
        };

        DexRewriter rewriter = new DexRewriter(module);

        List<ClassDef> modifiedClasses = new ArrayList<>(targetList.size());
        for (ClassDef classDef : targetList) {
            ClassDef rewrittenClass = rewriter.getClassDefRewriter().rewrite(classDef);
            modifiedClasses.add(ImmutableClassDef.of(rewrittenClass));
        }

        targetList.clear();
        targetList.addAll(modifiedClasses);
    }
}